package cn.me.alphamq.broker.rebalance;

import cn.me.alphamq.broker.cache.Container;
import cn.me.alphamq.broker.enumeration.RebalanceStrategyEnum;
import cn.me.alphamq.broker.rebalance.strategy.RebalanceStrategy;
import cn.me.alphamq.broker.rebalance.strategy.impl.RandomRebalance;
import cn.me.alphamq.broker.rebalance.strategy.impl.RoundRobinRebalance;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 消费者实例的管理者
 *
 * @author f
 */
public class ConsumerInstanceManager {

    private static final Logger log = LoggerFactory.getLogger(ConsumerInstanceManager.class);

    /**
     * 重平衡策略
     */
    private static RebalanceStrategy rebalanceStrategy;

    /**
     * 用于删除断开连接后的指定消费者的 map，key：ip:端口号，value：消费者实例
     */
    private static final Map<String, ConsumerInstance> CONSUMER_INSTANCE_MAP = new ConcurrentHashMap<>();

    static {
        RebalanceStrategyEnum rebalanceStrategy = Container.getBrokerProperties().getRebalanceStrategy();
        if (RebalanceStrategyEnum.ROUND_ROBIN == rebalanceStrategy) {
            ConsumerInstanceManager.rebalanceStrategy = new RoundRobinRebalance();
        } else if (RebalanceStrategyEnum.RANDOM == rebalanceStrategy) {
            ConsumerInstanceManager.rebalanceStrategy = new RandomRebalance();
        } else {
            throw new RuntimeException("rebalance strategy [" + rebalanceStrategy.getCode() + "] is not supported");
        }
    }

    /**
     * 根据消费者实例的参数获取消费者实例，如果消费者是第一次消费，则需要将消费者实例加入集合
     *
     * @param consumerInstance 封装消费者实例参数的对象
     * @return 需要的消费者实例
     */
    public static ConsumerInstance getInstance(ConsumerInstance consumerInstance) {
        String instanceId = getInstanceId(consumerInstance);
        Optional<ConsumerInstance> optional = Container.getConsumerInstanceMap()
                .computeIfAbsent(consumerInstance.getTopic(), k -> new HashMap<>())
                .computeIfAbsent(consumerInstance.getConsumerGroup(), k -> new ArrayList<>())
                .stream().filter(ci -> instanceId.equals(getInstanceId(ci))).findFirst();
        if (optional.isPresent()) {
            return optional.get();
        } else {
            addInstance(consumerInstance);
            return consumerInstance;
        }
    }

    /**
     * 重平衡有两个时机：
     * - 消费者实例的数量发生变化
     * - 队列的数量发生变化（目前的 MQ 不支持队列的数量发生变化）
     */
    private static void rebalance(String topic, String consumerGroup) {
        rebalanceStrategy.rebalance(topic, consumerGroup);
    }

    /**
     * 添加消费者实例
     *
     * @param instance 待添加的消费者实例
     */
    public static void addInstance(ConsumerInstance instance) {
        synchronized (ConsumerInstanceManager.class) {
            String topic = instance.getTopic();
            String consumerGroup = instance.getConsumerGroup();
            Container.getConsumerInstanceMap()
                    .get(topic)
                    .get(consumerGroup)
                    .add(instance);

            CONSUMER_INSTANCE_MAP.put(getInstanceId(instance), instance);
            rebalance(topic, consumerGroup);
        }
    }

    /**
     * 移除断开连接的消费者实例
     *
     * @param ip   消费者的 ip 地址
     * @param port 消费者的 端口号
     */
    public static void removeInactiveInstance(String ip, Integer port) {
        String instanceId = ip + ":" + port;
        if (CONSUMER_INSTANCE_MAP.containsKey(instanceId)) {
            removeInstance(CONSUMER_INSTANCE_MAP.remove(instanceId));
            log.info("remove consumer instance [{}]", instanceId);
        }
    }

    /**
     * 移除消费者实例
     *
     * @param instance 待移除的消费者实例
     */
    public static void removeInstance(ConsumerInstance instance) {
        synchronized (ConsumerInstanceManager.class) {
            String instanceId = getInstanceId(instance);
            String topic = instance.getTopic();
            String consumerGroup = instance.getConsumerGroup();
            List<ConsumerInstance> instanceList = Container.getConsumerInstanceMap()
                    .get(topic)
                    .get(consumerGroup);
            for (ConsumerInstance consumerInstance : instanceList) {
                if (getInstanceId(consumerInstance).equals(instanceId)) {
                    instanceList.remove(consumerInstance);
                    return;
                }
            }

            rebalance(topic, consumerGroup);
        }
    }

    /**
     * 获取消费者的唯一 id
     *
     * @param instance 待获取唯一 id 的消费者实例
     * @return 唯一 id
     */
    public static String getInstanceId(ConsumerInstance instance) {
        return getInstanceId(instance.getConsumerIp(), instance.getConsumerPort());
    }

    /**
     * 获取消费者的唯一 id
     *
     * @param ip   ip 地址
     * @param port 端口号
     * @return 唯一 id
     */
    public static String getInstanceId(String ip, Integer port) {
        return ip + ":" + port;
    }

}
